"""
推理相关工具函数
"""
import os
import time
import torch
import numpy as np
from PIL import Image
import cv2


def load_model_for_inference(model, checkpoint_path, device):
    """
    加载模型用于推理

    Args:
        model: 模型对象
        checkpoint_path: 检查点路径
        device: 设备

    Returns:
        model: 加载好权重的模型
    """
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"模型权重文件不存在: {checkpoint_path}")

    print(f"[INFO] 加载模型权重: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location=device)

    # 兼容不同的保存格式
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)

    model.to(device)
    model.eval()

    return model


def preprocess_image(image_path, to_rgb=False):
    """
    预处理输入图像

    Args:
        image_path: 图像路径
        to_rgb: 是否转换为3通道

    Returns:
        image_tensor: 预处理后的tensor
        original_image: 原始图像（用于可视化）
    """
    # 读取图像（灰度）
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if image is None:
        raise ValueError(f"无法读取图像: {image_path}")

    # 确保尺寸为256x256
    if image.shape != (256, 256):
        image = cv2.resize(image, (256, 256))

    # 归一化到[0, 1]
    image = image.astype(np.float32) / 255.0

    # 转换为tensor
    image_tensor = torch.from_numpy(image).unsqueeze(0)  # (1, H, W)

    # 如果需要，复制到3通道
    if to_rgb:
        image_tensor = image_tensor.repeat(3, 1, 1)  # (3, H, W)

    # 添加batch维度
    image_tensor = image_tensor.unsqueeze(0)  # (1, C, H, W)

    return image_tensor, image


def postprocess_prediction(pred, threshold=0.5):
    """
    后处理预测结果

    Args:
        pred: 模型输出的预测 (1, 1, H, W) 或 (1, 2, H, W)
        threshold: 二值化阈值

    Returns:
        mask: 二值化的mask (H, W)，值为0或255
    """
    # 如果是多通道输出，取前景通道
    if pred.shape[1] == 2:
        pred = torch.softmax(pred, dim=1)[:, 1:2, :, :]  # 取前景类
    else:
        pred = torch.sigmoid(pred)

    # 转换为numpy
    pred = pred.squeeze().cpu().numpy()

    # 二值化
    mask = (pred > threshold).astype(np.uint8) * 255

    return mask


def save_prediction(mask, output_path):
    """
    保存预测结果

    Args:
        mask: 预测的mask (H, W)，值为0或255
        output_path: 输出路径
    """
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    cv2.imwrite(output_path, mask)
    print(f"[INFO] 预测结果已保存: {output_path}")


def inference_single_image(model, image_path, output_path, device, to_rgb=False, threshold=0.5):
    """
    对单张图像进行推理

    Args:
        model: 模型
        image_path: 输入图像路径
        output_path: 输出路径
        device: 设备
        to_rgb: 是否转换为3通道
        threshold: 二值化阈值

    Returns:
        inference_time: 推理时间（秒）
        foreground_pixels: 前景像素数
    """
    # 预处理
    image_tensor, original_image = preprocess_image(image_path, to_rgb=to_rgb)
    image_tensor = image_tensor.to(device)

    # 推理
    start_time = time.time()
    with torch.no_grad():
        pred = model(image_tensor)
    inference_time = time.time() - start_time

    # 后处理
    mask = postprocess_prediction(pred, threshold=threshold)

    # 统计前景像素数
    foreground_pixels = np.sum(mask > 0)

    # 保存
    save_prediction(mask, output_path)

    return inference_time, foreground_pixels
